#include "scullc.h"

#include <linux/module.h>
#include <linux/init.h>
#include <linux/fs.h>
#include <linux/uaccess.h>
#include <linux/slab.h>
#include <linux/vmalloc.h>
#include <linux/device.h>
#include <linux/cdev.h>

struct scullc_qset {
    void *data;
    struct scullc_qset *next;
};

struct scullc_operations {
    void *(*zalloc)(size_t size, int flags);
    void (*free)(const void *ptr, size_t size);
};

struct scullc_dev {
    struct cdev cdev;
    struct scullc_qset *data;
    size_t size;
    struct scullc_operations *ops;
};

dev_t dev;
struct class *scullc_class;
struct scullc_dev *scullc_dev;
struct kmem_cache *scullc_slab_cache;

int allocator = 0;
module_param(allocator, int, S_IRUGO);
MODULE_PARM_DESC(allocator, "Memory allocator: 0=kmalloc (default), 1=vmalloc, 2=slab, 3=gfp");

static struct scullc_qset *scullc_follow(struct scullc_dev *dev, int n)
{
    struct scullc_qset *qs = dev->data;

    if (!qs) {
        qs = dev->data = dev->ops->zalloc(sizeof(struct scullc_qset), GFP_KERNEL);
        if (!qs)
            return NULL;
        qs->data = dev->ops->zalloc(DEFAULT_QSET, GFP_KERNEL);
        if (!qs->data)
            return NULL;
    }

    while (n--) {
        if (!qs->next) {
            qs->next = dev->ops->zalloc(sizeof(struct scullc_qset), GFP_KERNEL);
            if (!qs->next)
                return NULL;
            qs->next->data = dev->ops->zalloc(DEFAULT_QSET, GFP_KERNEL);
            if (!qs->next->data)
                return NULL;
        }
        qs = qs->next;
    }
    return qs;
}

static struct scullc_qset *scullc_lookup(struct scullc_dev *dev, int index)
{
    struct scullc_qset *qs = dev->data;

    while (index-- && qs) {
        qs = qs->next;
    }
    return qs;
}

static int scullc_open(struct inode *inode, struct file *filp)
{
    pr_info("scullc: Device opened\n");
    struct cdev *cdev = inode->i_cdev;
    struct scullc_dev *sdev = container_of(cdev, struct scullc_dev, cdev);

    filp->private_data = sdev;
    return 0;
}

static ssize_t scullc_read(struct file *filp, char __user *buf, size_t count, loff_t *f_pos)
{
    pr_info("scullc: Read %zu bytes from position %lld\n", count, *f_pos);
    struct scullc_dev *sdev = filp->private_data;
    struct scullc_qset *qset;
    int qset_index;
    size_t n = 0;

    pr_info("scullc: Device size %zu\n", sdev->size);
    if (*f_pos >= sdev->size) {
        return 0; // EOF
    }
    if (*f_pos + count > sdev->size) {
        count = sdev->size - *f_pos;
    }

    qset_index = *f_pos / DEFAULT_QSET;
    qset = scullc_lookup(sdev, qset_index);
    pr_info("scullc: Located qset index %d qset 0x%p\n", qset_index, qset);
    while (count > 0) {
        if (!qset || !qset->data) {
            break; // No more data
        }

        size_t offset = *f_pos % DEFAULT_QSET;
        size_t to_read = min(count, (size_t)(DEFAULT_QSET - offset));
        pr_info("scullc: Reading from qset index %d offset %zd count %zd\n", qset_index, offset, to_read);

        if (copy_to_user(buf, qset->data + offset, to_read)) {
            return -EFAULT;
        }

        buf += to_read;
        *f_pos += to_read;
        count -= to_read;
        n += to_read;
        qset = qset->next;
        qset_index++;
    }

    pr_info("scullc: Total bytes read %zu\n", n);
    return n;
}

static ssize_t scullc_write(struct file *filp, const char __user *buf, size_t count, loff_t *f_pos)
{
    pr_info("scullc: Write %zu bytes to position %lld\n", count, *f_pos);
    struct scullc_dev *sdev = filp->private_data;
    struct scullc_qset *qset;
    int qset_index;
    size_t n = 0;

    pr_info("scullc: Device size %zu\n", sdev->size);
    while (count > 0) {
        qset_index = *f_pos / DEFAULT_QSET;
        qset = scullc_follow(sdev, qset_index);
        if (!qset) {
            return -ENOMEM;
        }
        if (!qset->data) {
            return -ENOMEM;
        }

        size_t offset = *f_pos % DEFAULT_QSET;
        size_t to_write = min(count, (size_t)(DEFAULT_QSET - offset));
        pr_info("scullc: Writing to qset index %d offset %zd count %zd\n", qset_index, offset, to_write);

        if (copy_from_user(qset->data + offset, buf, to_write)) {
            return -EFAULT;
        }

        buf += to_write;
        *f_pos += to_write;
        count -= to_write;
        n += to_write;

        if (*f_pos > sdev->size) {
            sdev->size = *f_pos;
        }
    }

    pr_info("scullc: Total bytes written %zu total %zu\n", n, sdev->size);
    return n;
}

static int scullc_release(struct inode *inode, struct file *filp)
{
    pr_info("scullc: Device closed\n");
    return 0;
}

static struct file_operations scullc_fops = {
    .owner = THIS_MODULE,
    .open = scullc_open,
    .read = scullc_read,
    .write = scullc_write,
    .release = scullc_release,
};

static void *scullc_kzalloc(size_t size, int flags)
{
    return kzalloc(size, flags);
}

static void scullc_kfree(const void *ptr, size_t size)
{
    kfree(ptr);
}

static struct scullc_operations scullc_kmalloc_ops = {
    .zalloc = scullc_kzalloc,
    .free = scullc_kfree,
};

static void *scullc_vmalloc(size_t size, int flags)
{
    return vmalloc(size);
}

static void scullc_vfree(const void *ptr, size_t size)
{
    vfree(ptr);
}

static struct scullc_operations scullc_vmalloc_ops = {
    .zalloc = scullc_vmalloc,
    .free = scullc_vfree,
};

static void *scullc_slab_zalloc(size_t size, int flags)
{
    return kmem_cache_zalloc(scullc_slab_cache, flags);
}

static void scullc_slab_free(const void *ptr, size_t size)
{
    kmem_cache_free(scullc_slab_cache, (void *)ptr);
}

static struct scullc_operations scullc_slab_ops = {
    .zalloc = scullc_slab_zalloc,
    .free = scullc_slab_free,
};

static void *scullc_gfp_alloc(size_t size, int flags)
{
    if (size <= PAGE_SIZE) 
        return (void*)get_zeroed_page(flags);
    int order = get_order(size);
    return (void *)__get_free_pages(flags | __GFP_ZERO, order);
}

static void scullc_gfp_free(const void *ptr, size_t size)
{
    free_pages((unsigned long)ptr, get_order(size));
}

static struct scullc_operations scullc_gfp_ops = {
    .zalloc = scullc_gfp_alloc,
    .free = scullc_gfp_free,
};

static int scullc_setup(struct scullc_dev *dev)
{
    switch (allocator) {
    case 0:
        dev->ops = &scullc_kmalloc_ops;
        pr_info("scullc: Using kmalloc allocator\n");
        break;
    case 1:
        pr_info("scullc: Using vmalloc allocator\n");
        dev->ops = &scullc_vmalloc_ops;
        break;
    case 2:
        pr_info("scullc: Using slab allocator\n");
        scullc_slab_cache = kmem_cache_create("scullc_slab_cache", DEFAULT_QSET, 0, SLAB_HWCACHE_ALIGN, NULL);
        if (!scullc_slab_cache) {
            pr_err("scullc: Failed to create slab cache, defaulting to kmalloc\n");
            dev->ops = &scullc_kmalloc_ops;
            break;
        }
        dev->ops = &scullc_slab_ops;
        break;
    case 3:
        pr_info("scullc: Using gfp allocator\n");
        dev->ops = &scullc_gfp_ops;
        break;
    default:
        dev->ops = &scullc_kmalloc_ops;
        pr_info("scullc: Unknown allocator %d, defaulting to kmalloc\n", allocator);
        break;
    }
    return 0;
}

static __init int scullc_init(void)
{
    pr_info("scullc: Initialization started\n");
    struct device *device;
    int result;

    // Allocate device numbers
    result = alloc_chrdev_region(&dev, 0, 1, SCULLC_NAME);
    if (result < 0) {
        pr_err("scullc: Failed to allocate char device region\n");
        return result;
    }
    pr_info("scullc: Allocated device number %d:%d\n", MAJOR(dev), MINOR(dev));

    scullc_class = class_create(SCULLC_CLASS_NAME);
    if (IS_ERR(scullc_class)) {
        pr_err("scullc: Failed to create class\n");
        result = PTR_ERR(scullc_class);
        goto add_cdev_err;
    }
    pr_info("scullc: Created device class\n");

    scullc_dev = kzalloc(sizeof(struct scullc_dev), GFP_KERNEL);
    if (!scullc_dev) {
        pr_err("scullc: Failed to allocate cdev\n");
        result = -ENOMEM;
        goto alloc_cdev_err;
    }
    cdev_init(&scullc_dev->cdev, &scullc_fops);
    scullc_dev->cdev.owner = THIS_MODULE;

    scullc_setup(scullc_dev);

    result = cdev_add(&scullc_dev->cdev, dev, 1);
    if (result) {
        pr_err("scullc: Failed to add cdev\n");
        goto add_cdev_err;
    }
    device = device_create(scullc_class, NULL, dev, NULL, SCULLC_NAME "%d", 0);
    if (IS_ERR(device)) {
        pr_err("scullc: Failed to create device\n");
        result = PTR_ERR(device);
        goto add_cdev_err;
    }
    pr_info("scullc: Device created successfully\n");
    pr_info("scullc: Initialization completed\n");
    return 0;

alloc_cdev_err:
    unregister_chrdev_region(dev, 1);
add_cdev_err:
    kfree(scullc_dev);
    pr_info("scullc: Initialization failed\n");
    return result;
}
module_init(scullc_init);

static __exit void scullc_exit(void)
{
    pr_info("scullc: Cleanup started\n");
    struct scullc_qset *qs = scullc_dev->data;

    while (qs) {
        struct scullc_qset *next = qs->next;
        scullc_dev->ops->free(qs->data, DEFAULT_QSET);
        scullc_dev->ops->free(qs, DEFAULT_QSET);
        qs = next;
    }
    scullc_dev->data = NULL;
    scullc_dev->size = 0;
    if (scullc_slab_cache) {
        kmem_cache_destroy(scullc_slab_cache);
        scullc_slab_cache = NULL;
    }

    device_destroy(scullc_class, dev);
    class_destroy(scullc_class);
    cdev_del(&scullc_dev->cdev);
    kfree(scullc_dev);
    unregister_chrdev_region(dev, 1);
    pr_info("scullc: Cleanup completed\n");
}
module_exit(scullc_exit);

MODULE_LICENSE("GPL");